import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

# 读取数据
df = pd.read_excel('./associationRules/餐厅数据.xlsx')
print(df.head())

# 将数据转换为列表形式
trsations = df['菜品'].str.split(',').tolist()

# 编码交易数据
te = TransactionEncoder()
te_ary = te.fit(trsations).transform(trsations)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用Apriori算法找到频繁项集
frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 打印长度为2的频繁项集
print(frequent_itemsets[frequent_itemsets['itemsets'].apply(lambda x: len(x) == 2)])

# 生成关联规则
rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.1)

# 筛选有效规则
effective = rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
]
effective.sort_values(by=['lift', 'confidence'], ascending=False, inplace=True)
print(effective[['antecedents','consequents','support','confidence','lift']])

# 设置字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 关联关系可视化
G = nx.DiGraph()
for _, row in effective.iterrows():
    G.add_edge(':'.join(list(row['antecedents'])),
               ':'.join(list(row['consequents'])),
               weight=row['lift'])

plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G)
edge_colors = [G[u][v]['weight'] for u, v in G.edges()]
nx.draw(G, pos, with_labels=True, edge_color=edge_colors, node_size=500, font_size=10, font_weight='bold', node_color='skyblue', width=2.0)
plt.show()